package com.hjy.backend.indexManager;


import com.hjy.backend.commons.SubArray;
import com.hjy.backend.dataManager.DataManager;
import com.hjy.backend.dataManager.dataItem.DataItem;
import com.hjy.backend.transactionManager.TransactionManagerImpl;
import com.hjy.backend.utils.Parser;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * B+tree视频地址: https://www.bilibili.com/video/BV1HX4y1N7v6/
 */
public class BPlusTree {
    DataManager dm;
    long bootUid;
    DataItem bootDataItem; // 这个DataItem 中存储了UID , 8 位
    Lock bootLock;

    /**
     * 创建一颗 B+Tree
     *
     * @param dm
     * @return 注意, 这里新建了一个新的DataItem存储 rootUID
     * @throws Exception
     */
    public static long create(DataManager dm) throws Exception {
        byte[] rawRoot = Node.newNilRootRaw();
        // 使用超级事务插入节点到数据文件
        long rootUid = dm.insert(TransactionManagerImpl.SUPER_XID, rawRoot);
        // TODO ?
        return dm.insert(TransactionManagerImpl.SUPER_XID, Parser.long2Byte(rootUid));
    }

    /**
     * 根据根节点的 uid , 加载一颗B+tree
     *
     * @param bootUid
     * @param dm
     * @return
     * @throws Exception
     */
    public static BPlusTree load(long bootUid, DataManager dm) throws Exception {
        DataItem bootDataItem = dm.read(bootUid);
        assert bootDataItem != null;
        BPlusTree t = new BPlusTree();
        t.bootUid = bootUid;
        t.dm = dm;
        t.bootDataItem = bootDataItem;
        t.bootLock = new ReentrantLock();
        return t;
    }

    /**
     * 解析根节点的 uid
     *
     * @return
     */
    private long rootUid() {
        bootLock.lock();
        try {
            SubArray sa = bootDataItem.data();
            // 获取 UID, 占据 8 位
            return Parser.parseLong(Arrays.copyOfRange(sa.raw, sa.start, sa.start + 8));
        } finally {
            bootLock.unlock();
        }
    }

    /**
     * 根节点满了, 创建一个新的根节点, 开始根节点是叶子节点, 但后面随着索引的增多, 根节点就不是叶子节点了
     *
     * @param left 第一个节点的uid
     * @param right 第二个节点的uid
     * @param rightKey 新节点对应的第一个数据项的 key
     * @throws Exception
     */
    private void updateRootUid(long left, long right, long rightKey) throws Exception {
        bootLock.lock();
        try {
            byte[] rootRaw = Node.newRootRaw(left, right, rightKey);
            // 新根节点对应的 UID
            long newRootUid = dm.insert(TransactionManagerImpl.SUPER_XID, rootRaw);
            bootDataItem.before();
            SubArray diRaw = bootDataItem.data();
            // 更新根节点的uid
            System.arraycopy(Parser.long2Byte(newRootUid), 0, diRaw.raw, diRaw.start, 8);
            bootDataItem.after(TransactionManagerImpl.SUPER_XID);
        } finally {
            bootLock.unlock();
        }
    }

    /**
     * 根据 uid 和 key 寻找到对应数据项UKey 的叶子节点的 Uid
     *
     * @param nodeUid
     * @param key
     * @return
     * @throws Exception
     */
    private long searchLeaf(long nodeUid, long key) throws Exception {
        Node node = Node.loadNode(this, nodeUid);
        boolean isLeaf = node.isLeaf();
        node.release();

        if (isLeaf) {
            return nodeUid;
        } else {
            long next = searchNext(nodeUid, key);
            return searchLeaf(next, key);
        }
    }

    /**
     * 根据 ukey , 在同级树上从左到右寻找, 如果找到key所在的节点, 返回其存储的 uid
     * <p>
     * 如果所在的节点不是叶子节点, 返回的uid是下一级节点的uid
     * <p>
     * 如果是叶子节点, 返回的uid是对应记录的uid
     *
     * @param nodeUid
     * @param key
     * @return
     * @throws Exception
     */
    private long searchNext(long nodeUid, long key) throws Exception {
        while (true) {
            Node node = Node.loadNode(this, nodeUid);
            Node.SearchNextRes res = node.searchNext(key);
            node.release();
            if (res.uid != 0) return res.uid;
            nodeUid = res.siblingUid;
        }
    }

    /**
     * 查询key在哪个数据项
     *
     * @param key
     * @return 数据项对应的UID
     * @throws Exception
     */
    public List<Long> search(long key) throws Exception {
        return searchRange(key, key);
    }

    /**
     * 范围查找, 获取 leftKey 和 rightKey 之间所有的UID, 不管其是不是在不同的节点之间
     *
     * @param leftKey
     * @param rightKey
     * @return
     * @throws Exception
     */
    public List<Long> searchRange(long leftKey, long rightKey) throws Exception {
        long rootUid = rootUid();
        // 找到 leftKey 所在的叶子节点 uid
        long leafUid = searchLeaf(rootUid, leftKey);
        List<Long> uids = new ArrayList<>();
        while (true) {
            Node leaf = Node.loadNode(this, leafUid);
            // 在叶子节点那一层开始搜
            Node.LeafSearchRangeRes res = leaf.leafSearchRange(leftKey, rightKey);
            leaf.release();
            uids.addAll(res.uids);
            if (res.siblingUid == 0) {
                break;
            } else {
                leafUid = res.siblingUid;
            }
        }
        return uids;
    }

    /**
     * 从根节点开始搜, 找到合适的位置插入数据项(key)
     *
     * @param key
     * @param uid
     * @throws Exception
     */
    public void insert(long key, long uid) throws Exception {
        long rootUid = rootUid();
        // 从根节点开始寻找一个合适的位置进行插入
        InsertRes res = insert(rootUid, uid, key);
        assert res != null;
        // 根节点满了, 分裂了, 就需要一个新的根节点了
        if (res.newNode != 0) {
            updateRootUid(rootUid, res.newNode, res.newKey);
        }
    }

    class InsertRes {
        long newNode, newKey; // 新节点UID, 新节点第一个key
    }

    /**
     * 插入一个数据项
     * 叶子节点直接插入
     * 非叶子节点递归插入, 直到找到叶子节点再插入,
     *
     * @param nodeUid
     * @param uid
     * @param key
     * @return
     * @throws Exception
     */
    private InsertRes insert(long nodeUid, long uid, long key) throws Exception {
        Node node = Node.loadNode(this, nodeUid);
        boolean isLeaf = node.isLeaf();
        node.release();

        InsertRes res = null;
        if (isLeaf) {
            // 叶子节点直接插入
            res = insertAndSplit(nodeUid, uid, key);
        } else {
            // 非叶子节点递归插入, 直到找到叶子节点再插入
            long next = searchNext(nodeUid, key);
            InsertRes ir = insert(next, uid, key);
            if (ir.newNode != 0) {
                // 如果分割了节点, 这些非叶子节点就保存了一个新的数据项来访问到新的节点
                res = insertAndSplit(nodeUid, ir.newNode, ir.newKey);
            } else {
                res = new InsertRes();
            }
        }
        return res;
    }

    /**
     * 寻找到合适位置, 插入数据项
     *
     * @param nodeUid
     * @param uid
     * @param key
     * @return
     * @throws Exception
     */
    private InsertRes insertAndSplit(long nodeUid, long uid, long key) throws Exception {
        while (true) {
            Node node = Node.loadNode(this, nodeUid);
            // 插入数据项
            Node.InsertAndSplitRes iasr = node.insertAndSplit(uid, key);
            node.release();
            // 那就是插入失败咯, 因为key需要有序, 我们去下一个节点试试运气
            if (iasr.siblingUid != 0) {
                nodeUid = iasr.siblingUid;
            } else {
                // 插入成功
                InsertRes res = new InsertRes();
                res.newNode = iasr.newSon;
                res.newKey = iasr.newKey;
                return res;
            }
        }
    }

    public void close() {
        bootDataItem.release();
    }
}
